package prod_cons;

import java.util.concurrent.atomic.AtomicReference;

import prod_cons.TaskContainer;

public class TreesContainer implements TaskContainer {
	private static int treeHeight;
	private static int TREES_NUM = 2;

	static {
		String treesPerNodeStr = System.getProperty("treesPerNode");
		if (treesPerNodeStr != null) {
			TREES_NUM = Integer.parseInt(treesPerNodeStr);
			System.out.println("Using " + TREES_NUM + " trees per node");
		}
	}
	
	private static class NodeWithTree {
		private int id;
		
		public int getId() {
			return this.id;
		}

		public void setId(int id) {
			this.id = id;
		}

		public NodeWithTree() {
			for(int i = 0; i < trees.length; i++)
				trees[i] = new ScalableTreeContainer<Runnable>(treeHeight);
		}

		public ScalableTreeContainer[] trees = new ScalableTreeContainer[TREES_NUM];
		public AtomicReference<NodeWithTree> next = new AtomicReference<NodeWithTree>(null);
	}
	
	private static class ReferencePair {
		public NodeWithTree prev;
		public NodeWithTree cur;
		public ReferencePair(NodeWithTree prev, NodeWithTree cur) {
			this.prev = prev;
			this.cur = cur;
		}
	}

	private AtomicReference<NodeWithTree> PT = new AtomicReference<NodeWithTree>();
	private AtomicReference<ReferencePair> GT = new AtomicReference<ReferencePair>();

	public TreesContainer() {
		treeHeight = 10;
		String treeHeightStr = System.getProperty("initialHeight");
		if (treeHeightStr != null)
			treeHeight = Integer.parseInt(treeHeightStr);

		NodeWithTree firstNode = new NodeWithTree();
		firstNode.setId(0);
		PT.set(firstNode);
		GT.set(new ReferencePair(null, firstNode));
	}
	
	private static final ThreadLocal<NodeWithTree> localTree = new ThreadLocal<NodeWithTree>() {
		protected NodeWithTree initialValue() {
			return null;
		}
	};


	@Override
	public void put(Runnable task) {
		int treeId = ((int)Thread.currentThread().getId()) % TREES_NUM;
		while(true) {
			NodeWithTree latest = PT.get();
			if (latest.trees[treeId].put(task) == true) {
				ReferencePair newGT = null;
				while(true) {
					ReferencePair gtVal = GT.get();
					if (treeHeight <= 2 || gtVal.cur.id <= latest.id) return;
					if (newGT == null) newGT = new ReferencePair(null, latest);
					if (GT.compareAndSet(gtVal, newGT)) {
						return;
					}
				}
			} else {
				if (latest.next.get() == null) {
					insertNewTree();
				}
				PT.compareAndSet(latest, latest.next.get());
			}
		}
	}

	private void insertNewTree() {
		NodeWithTree newNode = localTree.get(); 
		if (newNode == null) newNode = new NodeWithTree();
		
		NodeWithTree cur = PT.get();
		newNode.setId(cur.getId() + 1);
		if (cur.next.compareAndSet(null, newNode)) {
			localTree.set(null);
		} else {
			localTree.set(newNode);
		}
	}

	@Override
	public Runnable get() {
		while(true) {
			ReferencePair gtVal = GT.get();
			Runnable task = null;
			
			if (gtVal.prev != null) task = getFromTrees(gtVal.prev.trees);
			if (task != null) { return task;}
			task = getFromTrees(gtVal.cur.trees);
			if (task != null) { return task;}
			
			// could not find a task in the tree
			if (PT.get().id > gtVal.cur.id) {
				ReferencePair newGT = new ReferencePair(gtVal.cur, gtVal.cur.next.get());
				GT.compareAndSet(gtVal, newGT);
			} else {
				return null;
			}
		}
	}
	
	private Runnable getFromTrees(ScalableTreeContainer[] trees) {
		Runnable task = null;
		for(ScalableTreeContainer<Runnable> tree : trees) {
			if ((task = tree.get()) != null) {
				return task;
			}
		}
		return null;
	}

	@Override
	public void setConsumersNum(int consumersNum) {
	}

	@Override
	public void registerConsumer() {
	}
}
